package net.lab1024.sa.ai.stream.module.business.ai.integration;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import lombok.extern.slf4j.Slf4j;
import net.lab1024.sa.ai.stream.module.business.ai.entity.AiModelConfigEntity;
import net.lab1024.sa.base.common.exception.BusinessException;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.*;
import org.springframework.stereotype.Service;
import org.springframework.web.client.HttpClientErrorException;
import org.springframework.web.client.RestTemplate;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Flux;
import java.util.*;

@Slf4j
@Service
public class ScnetAiService implements AiModelAdapter {
    private static final String MODEL_NAME = "scnet"; 
    @Autowired
    private LlmErrorCodeHandlerFactory errorCodeHandlerFactory;

    @Override
    public String chat(String prompt, Map<String, Object> params, AiModelConfigEntity config) {
        try {
            Map<String, Object> requestBody = new HashMap<>();
            requestBody.put("model", config.getModelType() != null ? config.getModelType() : "DeepSeek-R1-Distill-Qwen-7B");
            List<Map<String, String>> messages = new ArrayList<>();
            if (params != null && params.containsKey("messages")) {
                messages = (List<Map<String, String>>) params.get("messages");
            } else {
                messages.add(Map.of("role", "user", "content", prompt));
            }
            requestBody.put("messages", messages);
            requestBody.put("stream", false);
            if (params != null) {
                requestBody.putAll(params);
            }
            System.out.println("[Scnet] Request Body: " + JSON.toJSONString(requestBody));

            RestTemplate restTemplate = new RestTemplate();
            HttpHeaders headers = new HttpHeaders();
            headers.setContentType(MediaType.APPLICATION_JSON);
            if (config.getApiKey() != null && !config.getApiKey().isEmpty()) {
                headers.set("Authorization", "Bearer " + config.getApiKey());
            }
            HttpEntity<Map<String, Object>> entity = new HttpEntity<>(requestBody, headers);
            ResponseEntity<String> response = restTemplate.postForEntity(
                config.getEndpoint() != null ? config.getEndpoint() : "https://api.scnet.cn/api/llm/v1/chat/completions",
                entity, String.class);

            System.out.println("[Scnet] Response: " + response.getBody());

            JSONObject json = JSON.parseObject(response.getBody());
            if (json.containsKey("choices")) {
                return json.getJSONArray("choices").getJSONObject(0).getJSONObject("message").getString("content");
            } else {
                return response.getBody();
            }
        } catch (HttpClientErrorException e) {
            String msg = errorCodeHandlerFactory
                .getHandler(MODEL_NAME)
                .getFriendlyMessage(e.getStatusCode().value(), e.getResponseBodyAsString());
            // throw new BusinessException(msg);
            return msg;
        }
    }

    @Override
    public Flux<String> chatStream(String prompt, Map<String, Object> params, AiModelConfigEntity config) {
        log.info("[Scnet] chatStream called, prompt: {}, roleId: {}, promptIds: {}", prompt, params != null ? params.get("roleId") : null, params != null ? params.get("promptIds") : null);
        try {
            Map<String, Object> requestBody = new HashMap<>();
            requestBody.put("model", config.getModelType() != null ? config.getModelType() : "DeepSeek-R1-Distill-Qwen-7B");
            List<Map<String, String>> messages = new ArrayList<>();
            if (params != null && params.containsKey("messages")) {
                messages = (List<Map<String, String>>) params.get("messages");
            } else {
                messages.add(Map.of("role", "user", "content", prompt));
            }
            requestBody.put("messages", messages);
            requestBody.put("stream", true);
            if (params != null) {
                requestBody.putAll(params);
            }
            messages.stream().filter(m -> "system".equals(m.get("role"))).forEach(m -> log.info("[Scnet] system message: {}", m.get("content")));
            log.info("[Scnet] Request Body: {}", requestBody);

            WebClient webClient = WebClient.builder()
                .baseUrl(config.getEndpoint() != null ? config.getEndpoint() : "https://api.scnet.cn/api/llm/v1/chat/completions")
                .defaultHeader(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_JSON_VALUE)
                .defaultHeader("Authorization", config.getApiKey() != null ? "Bearer " + config.getApiKey() : "")
                .defaultHeader(HttpHeaders.ACCEPT, MediaType.ALL_VALUE)
                .build();

            return webClient.post()
                .bodyValue(requestBody)
                .exchangeToFlux(response -> {
                    String contentType = response.headers().contentType().map(MediaType::toString).orElse("");
                    if (contentType.contains("text/event-stream")) {
                        // SSE流式，直接解析每个chunk
                        return response.bodyToFlux(String.class)
                            .filter(chunk -> chunk != null && !chunk.trim().isEmpty())
                            .flatMap(chunk -> {
                                // log.info("[Scnet] SSE原始chunk: {}", chunk);
                                if ("[DONE]".equals(chunk.trim())) {
                                    log.info("[Scnet] Received [DONE], stream end");
                                    return Flux.empty();
                                }
                                try {
                                    JSONObject json = JSON.parseObject(chunk);
                                    if (json.containsKey("choices")) {
                                        JSONObject choice = json.getJSONArray("choices").getJSONObject(0);
                                        if (choice.containsKey("delta") && choice.getJSONObject("delta").containsKey("content")) {
                                            String content = choice.getJSONObject("delta").getString("content");
                                            // log.info("[Scnet] emit to client: {}", content);
                                            return content != null ? Flux.just(content) : Flux.empty();
                                        }
                                    }
                                } catch (Exception e) {
                                    log.warn("[Scnet] SSE chunk parse error", e);
                                }
                                return Flux.empty();
                            });
                    } else {
                        // 普通JSON响应
                        return response.bodyToMono(String.class)
                            .flatMapMany(body -> {
                                log.info("[Scnet] 普通JSON响应: {}", body);
                                try {
                                    JSONObject json = JSON.parseObject(body);
                                    if (json.containsKey("choices")) {
                                        JSONObject choice = json.getJSONArray("choices").getJSONObject(0);
                                        if (choice.containsKey("message") && choice.getJSONObject("message").containsKey("content")) {
                                            String content = choice.getJSONObject("message").getString("content");
                                            return Flux.just(content);
                                        }
                                    }
                                } catch (Exception e) {
                                    log.warn("[Scnet] 普通JSON解析失败", e);
                                }
                                return Flux.just(body);
                            });
                    }
                })
                .doOnError(e -> log.error("[Scnet] WebClient error", e))
                .onErrorResume(e -> {
                    String msg = "[AI服务连接失败] " + (e.getMessage() != null ? e.getMessage() : "未知错误");
                    log.error("[Scnet] onErrorResume: {}", msg, e);
                    return Flux.just(msg);
                });
        } catch (Exception e) {
            log.error("[Scnet] chatStream outer exception", e);
            return Flux.just("[AI服务异常] " + (e.getMessage() != null ? e.getMessage() : "未知错误"));
        }
    }

    @Override
    public String getModelName() {
        return MODEL_NAME;
    }
} 